# -*- encoding: utf-8 -*-
'''
@File    :   jiangnan.py
@Time    :   2022/2/11 15:20
@Author  :   ZhangChaoYang
@Desc    :   
'''

import os, sys

import numpy as np

sys.path.insert(0, os.getcwd())
from util.work_flow import trans_collection, transform
import re

regex = re.compile("[a-z]{1,2}(\d+)_.*")

fs = 50e3  # 采样频率50kHz
class2file = {'normal': 'n', 'ball': 'tb', 'inner': 'ib', 'outer': 'ob'}
WC = ["600", "800", "1000"]
corpus_dir = 'E:\\zcy\\数据集\\JiangNan'
corpus_name = "jiangnan"


def gen_corpus(trans, dim):
    if trans not in trans_collection:
        raise Exception("unsupported transformation method: {}".format(trans))
    for file in os.listdir(corpus_dir):
        if not file.endswith(".csv"):
            continue
        if file.startswith('n'):
            cls = 'normal'
        elif file.startswith('tb'):
            cls = 'ball'
        elif file.startswith('ib'):
            cls = 'inner'
        elif file.startswith('ob'):
            cls = 'outer'
        else:
            continue
        rpm = regex.findall(file)
        if len(rpm) == 0:
            return None
        rpm = rpm[0]
        infile = os.path.join(corpus_dir, file)
        fin = open(infile, encoding="utf-8")
        lines = fin.readlines()
        array = []
        for line in lines:
            content = line.strip()
            if len(content) > 0:
                array.append(float(content))

        x = transform(trans, array, 1024, 1.0, fs)
        print(cls, rpm, x.shape)
        out_path = os.path.join("corpus", "multi_class", dim, corpus_name, trans, rpm)
        if not os.path.exists(out_path):
            os.makedirs(out_path)
        np.save(os.path.join(out_path, cls), x)


if __name__ == '__main__':
    gen_corpus('original', '1d')
    gen_corpus('fft', '1d')
    gen_corpus('stat', '1d')
    gen_corpus('sfft', '2d')
    # gen_corpus('cwt', '2d')
